from keras.models import Sequential
from keras import layers

def set_model(dropout):
    model = Sequential()
    model.add(layers.Conv2D(16, (3,3), activation='relu', input_shape=
                            (image_height,image_width,NUM_CHANNELS)))
    if (dropout):
        model.add(layers.Dropout(0.3))
    model.add(layers.MaxPooling2D((2,2)))
    model.add(layers.Conv2D(32, (3,3), activation='relu'))
    if (dropout):
        model.add(layers.Dropout(0.3))
    model.add(layers.MaxPooling2D((2,2)))
    model.add(layers.Conv2D(64, (3,3), activation='relu'))  
    model.add(layers.Flatten())
    model.add(layers.Dense(32, activation='relu')) 
    model.add(layers.Dense(num_classes, activation='softmax'))
    
    if (dropout):
        print(model.summary())
  
    return model

EPOCHS = 20
for dropout in [True,False]:
    model = set_model(dropout)
    model.compile(loss='categorical_crossentropy', optimizer='rmsprop',
                  metrics=['accuracy'])
    history = model.fit(X_train, y_train, validation_data=(X_test,
                        y_test), epochs=EPOCHS, batch_size=128, verbose=0)
    print("Dropout: " + str(dropout))
    plot(history)
